import pymysql
import json
from datetime import datetime

class Database:
    def __init__(self, host='localhost', user='root', password='123456', db='face_authentication', port=3306):
        self.host = host
        self.user = user
        self.password = password
        self.db = db
        self.port = port
        self.conn = None
        self.cursor = None
        # 先尝试创建数据库（如果不存在）
        self._create_database_if_not_exists()
        self.init_db()
    
    def _create_database_if_not_exists(self):
        """创建数据库（如果不存在）"""
        try:
            # 先不指定数据库连接
            temp_conn = pymysql.connect(
                host=self.host,
                user=self.user,
                password=self.password,
                port=self.port,
                charset='utf8mb4'
            )
            temp_cursor = temp_conn.cursor()
            
            # 创建数据库（如果不存在）
            temp_cursor.execute(f"CREATE DATABASE IF NOT EXISTS `{self.db}` DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci")
            temp_conn.commit()
            temp_cursor.close()
            temp_conn.close()
            print(f"数据库 {self.db} 检查/创建成功")
        except Exception as e:
            print(f"创建数据库时出错: {e}")
            # 这里不抛出异常，允许程序继续尝试连接，可能数据库已存在但有其他问题
    
    def connect(self):
        """连接到MySQL数据库"""
        try:
            self.conn = pymysql.connect(
                host=self.host,
                user=self.user,
                password=self.password,
                db=self.db,
                port=self.port,
                charset='utf8mb4',
                cursorclass=pymysql.cursors.DictCursor
            )
            self.cursor = self.conn.cursor()
        except Exception as e:
            print(f"数据库连接失败: {e}")
            raise
    
    def disconnect(self):
        """断开数据库连接"""
        if self.conn:
            self.conn.close()
            self.conn = None
            self.cursor = None
    
    def init_db(self):
        """初始化数据库，创建用户表"""
        self.connect()
        self.cursor.execute('''
        CREATE TABLE IF NOT EXISTS users (
            id INT AUTO_INCREMENT PRIMARY KEY,
            name VARCHAR(255) NOT NULL,
            face_encoding TEXT NOT NULL,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            last_login TIMESTAMP NULL
        )
        ''')
        self.conn.commit()
        self.disconnect()
    
    def add_user(self, name, face_encoding):
        """添加新用户及其人脸编码"""
        self.connect()
        try:
            # 检查是否需要转换为列表
            if hasattr(face_encoding, 'tolist'):
                # 如果是numpy数组，先转换为列表
                face_encoding_list = face_encoding.tolist()
            else:
                # 假设已经是列表格式
                face_encoding_list = face_encoding
            
            # 将列表转换为JSON字符串存储
            face_encoding_str = json.dumps(face_encoding_list)
            
            self.cursor.execute(
                "INSERT INTO users (name, face_encoding) VALUES (%s, %s)",
                (name, face_encoding_str)
            )
            user_id = self.cursor.lastrowid
            self.conn.commit()
            return user_id
        except Exception as e:
            print(f"添加用户时出错: {e}")
            self.conn.rollback()
            return None
        finally:
            self.disconnect()
    
    def get_all_users(self):
        """获取所有用户信息"""
        self.connect()
        self.cursor.execute("SELECT id, name, created_at, last_login FROM users")
        users = self.cursor.fetchall()
        self.disconnect()
        return users
    
    def get_user_by_id(self, user_id):
        """根据ID获取用户信息"""
        self.connect()
        self.cursor.execute(
            "SELECT id, name, face_encoding FROM users WHERE id = %s",
            (user_id,)
        )
        user = self.cursor.fetchone()
        self.disconnect()
        if user:
            user['face_encoding'] = json.loads(user['face_encoding'])
        return user
    
    def update_last_login(self, user_id):
        """更新用户最后登录时间"""
        self.connect()
        try:
            current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            self.cursor.execute(
                "UPDATE users SET last_login = %s WHERE id = %s",
                (current_time, user_id)
            )
            self.conn.commit()
        except Exception as e:
            print(f"更新最后登录时间出错: {e}")
            self.conn.rollback()
        finally:
            self.disconnect()
    
    def get_all_face_encodings(self):
        """获取所有用户的人脸编码"""
        self.connect()
        self.cursor.execute("SELECT id, name, face_encoding FROM users")
        users = self.cursor.fetchall()
        self.disconnect()
        # 解析JSON字符串为Python对象
        for user in users:
            user['face_encoding'] = json.loads(user['face_encoding'])
        return users
    
    def delete_user(self, user_id):
        """删除用户"""
        self.connect()
        try:
            self.cursor.execute("DELETE FROM users WHERE id = %s", (user_id,))
            affected_rows = self.cursor.rowcount
            self.conn.commit()
            return affected_rows > 0
        except Exception as e:
            print(f"删除用户时出错: {e}")
            self.conn.rollback()
            return False
        finally:
            self.disconnect()